from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from PIL import Image
import io
import os
import cv2
import numpy as np
from rapid_table import RapidTable
import pandas as pd
import traceback
import logging

# 配置日志
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# 初始化RapidTable
table_engine = RapidTable()

def compress_png(img, quality):
    """PNG特定的压缩策略"""
    output = io.BytesIO()
    
    # 如果质量为100，只进行格式转换，不做任何压缩
    if quality >= 100:
        img.save(output, 
                format='PNG',
                optimize=False)  # 不进行额外优化
        return output
    
    # 转换为P模式(8位调色板)来减少颜色数量
    if quality < 85 and img.mode != 'P':
        colors = int(256 * quality / 100)  # 根据质量调整颜色数量
        img = img.convert('P', palette=Image.Palette.ADAPTIVE, colors=colors)
    
    # 缩放图片
    if quality < 50:
        width, height = img.size
        new_width = int(width * 0.8)
        new_height = int(height * 0.8)
        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 保存时使用最高压缩级别
    img.save(output, 
            format='PNG',
            optimize=True,
            compress_level=9)
    return output

def compress_jpeg(img, quality):
    """JPEG特定的压缩策略"""
    output = io.BytesIO()
    
    # 如果质量为100，只进行格式转换，不做任何压缩
    if quality >= 100:
        if img.mode != 'RGB':
            img = img.convert('RGB')
        img.save(output, 
                format='JPEG',
                quality=100,  # 使用最高质量
                optimize=False,  # 不进行额外优化
                progressive=False)  # 不使用渐进式
        return output
    
    # 缩放图片
    if quality < 50:
        width, height = img.size
        new_width = int(width * 0.8)
        new_height = int(height * 0.8)
        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 转换为RGB模式
    if img.mode != 'RGB':
        img = img.convert('RGB')
    
    # 保存时使用更激进的压缩参数
    img.save(output, 
            format='JPEG',
            quality=max(quality - 10, 10),  # 降低10%的质量，但不低于10%
            optimize=True,
            progressive=True,
            subsampling='4:2:0')
    return output

@app.route('/api/compress', methods=['POST'])
def compress_image():
    if 'image' not in request.files:
        return jsonify({'error': '没有上传图片'}), 400
    
    try:
        file = request.files['image']
        quality = int(request.form.get('quality', 85))
        output_format = request.form.get('format', 'auto')
        
        # 读取原始图片数据到内存
        input_data = file.read()
        input_stream = io.BytesIO(input_data)
        img = Image.open(input_stream)
        
        # 获取原始格式
        original_format = img.format or 'JPEG'
        
        # 如果指定了输出格式且不是auto，使用指定的格式
        if output_format != 'auto':
            original_format = output_format.upper()
        
        # 根据图片格式选择不同的压缩策略
        try:
            if original_format == 'PNG':
                output = compress_png(img, quality)
            else:  # 默认使用JPEG
                output = compress_jpeg(img, quality)
            
            output.seek(0)
            compressed_data = output.getvalue()
            compressed_size = len(compressed_data)
            
            # 如果压缩效果不理想，尝试另一种格式
            if compressed_size >= len(input_data):
                # PNG转JPEG可能会得到更好的压缩效果
                if original_format == 'PNG' and 'A' not in img.mode:  # 没有透明通道时才转换
                    output = compress_jpeg(img, quality)
                    output.seek(0)
                    compressed_data = output.getvalue()
                    compressed_size = len(compressed_data)
                    if compressed_size < len(input_data):
                        original_format = 'JPEG'
            
            # 如果压缩后仍然更大，返回原图
            if compressed_size >= len(input_data):
                return send_file(
                    io.BytesIO(input_data),
                    mimetype=f'image/{original_format.lower()}',
                    as_attachment=True,
                    download_name=f'original_{file.filename}'
                )
                
            return send_file(
                io.BytesIO(compressed_data),
                mimetype=f'image/{original_format.lower()}',
                as_attachment=True,
                download_name=f'compressed_{file.filename}'
            )
            
        except Exception as e:
            # 如果压缩过程出错，返回原图
            return send_file(
                io.BytesIO(input_data),
                mimetype=f'image/{original_format.lower()}',
                as_attachment=True,
                download_name=f'original_{file.filename}'
            )
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/recognize-table', methods=['POST'])
def recognize_table():
    if 'image' not in request.files:
        return jsonify({'error': '没有上传图片'}), 400
    
    try:
        file = request.files['image']
        # 读取图片数据
        file_bytes = file.read()
        nparr = np.frombuffer(file_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            logger.error('无法读取图片数据')
            return jsonify({'error': '无法读取图片'}), 400
        
        # 使用RapidTable识别表格
        try:
            # 转换图片格式
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            logger.debug(f'图片尺寸: {img_rgb.shape}')
            
            # 识别表格
            logger.debug('开始识别表格...')
            result = table_engine(img_rgb)
            
            # 根据返回值类型处理结果
            if isinstance(result, tuple):
                if len(result) == 2:
                    table_html, cells = result
                    if isinstance(cells, (list, tuple)):
                        logger.debug(f'识别结果: {len(cells)} 个单元格')
                    else:
                        logger.debug('识别结果: 无法获取单元格数量')
                else:
                    table_html = result[0]
                    logger.debug('识别结果: 无法获取单元格信息')
            else:
                table_html = result
                logger.debug('识别结果: 仅返回HTML')
            
            if not table_html:
                logger.warning('未识别到表格内容')
                return jsonify({'error': '未能识别到表格'}), 400
            
            # 确保table_html是字符串类型
            if not isinstance(table_html, str):
                table_html = str(table_html)
            
            return jsonify({
                'success': True,
                'html': table_html
            })
            
        except Exception as e:
            logger.error(f'表格识别失败: {str(e)}')
            logger.error(f'错误详情: {traceback.format_exc()}')
            return jsonify({'error': f'表格识别失败: {str(e)}'}), 500
    
    except Exception as e:
        logger.error(f'处理请求失败: {str(e)}')
        logger.error(f'错误详情: {traceback.format_exc()}')
        return jsonify({'error': str(e)}), 500

@app.route('/api/export-excel', methods=['POST'])
def export_excel():
    try:
        html = request.json['html']
        # 读取HTML表格并替换空值
        df = pd.read_html(html)[0]
        df = df.fillna('')  # 将 NaN 替换为空字符串
        
        output = io.BytesIO()
        
        # 使用 ExcelWriter 以便设置列宽
        with pd.ExcelWriter(output, engine='openpyxl') as writer:
            df.to_excel(writer, index=False, na_rep='')
            
            # 获取工作表
            worksheet = writer.sheets['Sheet1']
            
            # 调整每列的宽度
            for idx, col in enumerate(df.columns):
                # 计算列中最长内容的长度
                max_length = max(
                    df[col].astype(str).apply(len).max(),  # 数据的最大长度
                    len(str(col))  # 列名的长度
                )
                # 设置列宽（加上一些padding）
                worksheet.column_dimensions[chr(65 + idx)].width = max_length + 2
        
        output.seek(0)
        
        return send_file(
            output,
            mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            as_attachment=True,
            download_name='table.xlsx'
        )
    except Exception as e:
        logger.error(f'导出Excel失败: {str(e)}')
        logger.error(f'错误详情: {traceback.format_exc()}')
        return jsonify({'error': str(e)}), 500

@app.route('/api/export-word', methods=['POST'])
def export_word():
    try:
        html = request.json['html']
        from docx import Document
        from docx.shared import Inches
        import pandas as pd
        from bs4 import BeautifulSoup
        
        # 创建Word文档
        doc = Document()
        
        try:
            # 使用pandas读取HTML表格并替换空值
            df = pd.read_html(html)[0]
            df = df.fillna('')  # 将 NaN 替换为空字符串
            
            # 创建表格
            rows, cols = df.shape
            table = doc.add_table(rows=rows+1, cols=cols)  # +1 为表头行
            table.style = 'Table Grid'  # 添加表格边框
            
            # 添加表头
            for col_idx, column in enumerate(df.columns):
                cell = table.cell(0, col_idx)
                cell.text = str(column)
            
            # 添加数据
            for row_idx, row in enumerate(df.values):
                for col_idx, value in enumerate(row):
                    cell = table.cell(row_idx + 1, col_idx)  # +1 跳过表头行
                    # 将 nan 和 None 转换为空字符串
                    cell.text = '' if pd.isna(value) else str(value)
            
            # 自动调整列宽
            for cell in table.columns:
                cell.width = Inches(1.5)
            
        except Exception as e:
            logger.error(f'转换表格失败: {str(e)}')
            # 如果转换失败，尝试直接添加文本
            doc.add_paragraph(html)
        
        # 保存文档
        output = io.BytesIO()
        doc.save(output)
        output.seek(0)
        
        return send_file(
            output,
            mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document',
            as_attachment=True,
            download_name='table.docx'
        )
    except Exception as e:
        logger.error(f'导出Word失败: {str(e)}')
        logger.error(f'错误详情: {traceback.format_exc()}')
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True) 